from mindspore.train.callback import TimeMonitor
from utils import *
import argparse
import mindspore.dataset as ds
from models import *
import mindspore.nn as nn
from mindspore import Model
from mindspore import load_checkpoint, load_param_into_net
import mindspore.dataset.transforms.c_transforms as C
from mindspore import dtype as mstype
from mindspore.nn.metrics import Accuracy
parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str,  default='DPCNN', help='choose a model: TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer')
parser.add_argument('--embedding', default='random', type=str, help='random or pre_trained')
parser.add_argument('--word', default=True, type=bool, help='True for word, False for char')
args = parser.parse_args()


if __name__ =='__main__':
    dataset = '.'  # 文件名前缀
    embedding = 'random'
    model_name = args.model  # 'TextRCNN'  # TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer
    config = Config(dataset)
    np.random.seed(1)
    start_time = time.time()
    print("Loading data......")

# load data
    vocab, train_data, dev_data, test_data = build_dataset(config, True)

    # dev
    test_iter = DatasetMSIterater(test_data, config)
    test_data = ds.GeneratorDataset(test_iter, ['data', 'label'], num_parallel_workers=1, shuffle=True)
    test_dataset = test_data.map(C.TypeCast(mstype.int32), input_columns='label', num_parallel_workers=1)

# init net
    config.n_vocab = len(vocab)
    net = DPCNN(config)


# todo check if param updated
    param_dict = load_checkpoint("./model/DPCNN_1-2_1350.ckpt")
    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
    load_param_into_net(net, param_dict)
    model = Model(net, loss_fn = loss, metrics = {"Accuracy": Accuracy()})


# test
    print("============== Starting Inferencing ==============")
    acc = model.eval(test_dataset, dataset_sink_mode=False)
    print(acc['Accuracy']*100)
    print("============== Inferencing Success ==============")



